# 特征提取模块

import torch.nn as nn
import torch.nn.functional as F
import src.global_config as gl
import src.tool as tool


# 基础块?
class BasicBlock(nn.Module):
    def __init__(self, inplanes, planes, config_block, stride=1, downsample=None):
        self.BLOCKS = config_block
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)

        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = F.relu(self.bn1(out))

        out = self.conv2(out)
        out = self.bn2(out)

        out = out + residual
        out = F.relu(out)

        return out


# 提取器
class Extractor(nn.Module):
    # 17 10 10
    def __init__(self, inplanes, outplanes, config_block):
        super(Extractor, self).__init__()
        self.BLOCKS = config_block
        self.conv1 = nn.Conv2d(inplanes, outplanes, stride=1, kernel_size=3,
                               padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(outplanes)

        for block in range(self.BLOCKS):
            setattr(self, "res{}".format(block), BasicBlock(outplanes, outplanes, config_block))

    def forward(self, x):
        """
        x : tensor representing the state
        feature_maps : result of the residual layers forward pass
        """
        x = F.relu(self.bn1(self.conv1(x)))
        for block in range(self.BLOCKS - 1):
            x = getattr(self, "res{}".format(block))(x)

        feature_maps = getattr(self, "res{}".format(self.BLOCKS - 1))(x)
        return feature_maps
